from collections import deque
from typing import List, Tuple, Dict, Optional
import random
import torch
import numpy as np
import pandas as pd
import math


class ScheduleQAgent:
    """Q-learning agent for course scheduling optimization"""

    def __init__(self):
        self.q_table = {}
        self.learning_rate = 0.1
        self.discount_factor = 0.95
        self.epsilon = 0.1
        self.memory = deque(maxlen=2000)

        # Combined weight ranges from both versions
        self.weight_ranges = {
            'scv_weight': (0.2, 0.5),
            'weekend_weight': (0.15, 0.35),
            'efficient_weight': (0.15, 0.35),
            'priority_weight': (0.1, 0.3)
        }

        # Parameter ranges for scoring calculations
        self.parameter_ranges = {
            'scv_base': (50, 90),
            'deduct_factor': (5, 15),
            'reward_factor': (8, 20),
            'dense_factor': (10, 25),
            'reward_deviation': (5, 15),
            'weekend_base': (80, 120),
            'weekend_penalty': (40, 80),
            'efficient_base': (60, 100),
            'priority_base': (40, 80)
        }

    def get_state(self, schedule_tensor: torch.Tensor,
                  priority_data: pd.DataFrame) -> Tuple[float, int]:
        occupancy = schedule_tensor.sum().item()
        priority_level = int(priority_data['Application_Priority'].mode()[0])
        return (occupancy, priority_level)

    def _dict_to_tuple(self, d: Dict[str, float]) -> Tuple[Tuple[str, float], ...]:
        """Convert dictionary to hashable tuple of tuples"""
        return tuple(sorted(d.items()))

    def _tuple_to_dict(self, t: Tuple[Tuple[str, float], ...]) -> Dict[str, float]:
        """Convert tuple back to dictionary"""
        return dict(t)

    def select_weights(self, state: Tuple[float, int]) -> Dict[str, float]:
        if random.random() < self.epsilon:
            return {k: random.uniform(v[0], v[1])
                    for k, v in self.weight_ranges.items()}

        if state not in self.q_table:
            self.q_table[state] = {}

        if self.q_table[state]:
            weights_tuple = max(self.q_table[state].items(), key=lambda x: x[1])[0]
            return self._tuple_to_dict(weights_tuple)

        return {k: np.mean(v) for k, v in self.weight_ranges.items()}

    def update(self, state: Tuple[float, int], weights: Dict[str, float],
               reward: float, next_state: Tuple[float, int]) -> None:
        if state not in self.q_table:
            self.q_table[state] = {}

        weights_tuple = self._dict_to_tuple(weights)
        if weights_tuple not in self.q_table[state]:
            self.q_table[state][weights_tuple] = 0

        next_max = (max(self.q_table[next_state].values())
                    if next_state in self.q_table and self.q_table[next_state] else 0)

        self.q_table[state][weights_tuple] = (
                (1 - self.learning_rate) * self.q_table[state][weights_tuple] +
                self.learning_rate * (reward + self.discount_factor * next_max)
        )


# ... rest of the code remains the same ...


def calculate_course_density(day: int, schedule_tensor: torch.Tensor) -> float:
    """Calculate course density for a given day"""
    daily_schedule = schedule_tensor[:, day, :]
    total_courses = torch.sum(daily_schedule).item()
    return total_courses / 12


def calculate_standard_deviation(courses_per_week: List[int]) -> float:
    """Calculate standard deviation of course distribution"""
    avg = sum(courses_per_week[:5]) / 5  # Only consider weekdays
    squared_diffs = [(x - avg) ** 2 for x in courses_per_week[:5]]
    variance = sum(squared_diffs) / 4
    return math.sqrt(variance)


class ScheduleEvaluator:
    """Evaluates course schedules using multiple criteria"""

    def __init__(self, agent: ScheduleQAgent):
        self.agent = agent

    def evaluate_scv(self, schedule: torch.Tensor,
                     priority_data: pd.Series) -> float:
        point = self.agent.parameter_ranges['scv_base'][0]

        # Penalize for 12th period classes
        if priority_data['Course_Sequence'][-1] == "12":
            point -= self.agent.parameter_ranges['deduct_factor'][0]

        # Reward for optimal starting periods
        if int(priority_data['Course_Sequence'][0]) % 2 == 1:
            point += self.agent.parameter_ranges['reward_factor'][0]

        # Consider course density
        density = calculate_course_density(priority_data['Day'], schedule)
        point += density * self.agent.parameter_ranges['dense_factor'][0]

        # Consider schedule distribution
        courses_per_week = [5, 5, 5, 5, 5, 0, 0]  # Ideal distribution
        std_dev = calculate_standard_deviation(courses_per_week)
        point += std_dev * self.agent.parameter_ranges['reward_deviation'][0]

        return point

    def evaluate_weekend(self, schedule: torch.Tensor,
                         priority_data: pd.Series) -> float:
        base = self.agent.parameter_ranges['weekend_base'][0]
        day = priority_data['Day']

        if day in [6, 7]:  # Weekend penalty
            base -= self.agent.parameter_ranges['weekend_penalty'][0]

        return base

    def evaluate_efficiency(self, schedule: torch.Tensor,
                            priority_data: pd.Series) -> float:
        base = self.agent.parameter_ranges['efficient_base'][0]
        efficiency = (int(priority_data['Number_Of_Students']) /
                      int(priority_data['Volume']))
        return base * efficiency

    def evaluate_priority(self, schedule: torch.Tensor,
                          priority_data: pd.Series) -> float:
        base = self.agent.parameter_ranges['priority_base'][0]
        priority_level = priority_data['Application_Priority']

        priority_bonus = {
            "1": 50,
            "2": 30,
            "3": 10
        }

        return base + priority_bonus.get(priority_level, 0)


def judge_schedule(schedule_tensor: torch.Tensor,
                   priority_df: pd.DataFrame) -> List[float]:
    """Main function to evaluate course schedules"""
    agent = ScheduleQAgent()
    evaluator = ScheduleEvaluator(agent)
    results = []

    for _, row in priority_df.iterrows():
        state = agent.get_state(schedule_tensor, priority_df)
        weights = agent.select_weights(state)

        # Calculate individual scores
        scv = evaluator.evaluate_scv(schedule_tensor, row)
        weekend = evaluator.evaluate_weekend(schedule_tensor, row)
        efficient = evaluator.evaluate_efficiency(schedule_tensor, row)
        priority = evaluator.evaluate_priority(schedule_tensor, row)

        # Calculate weighted score
        score = (weights['scv_weight'] * scv +
                 weights['weekend_weight'] * weekend +
                 weights['efficient_weight'] * efficient +
                 weights['priority_weight'] * priority)

        results.append(round(score, 6))

        # Update Q-table
        next_state = agent.get_state(schedule_tensor, priority_df)
        agent.update(state, weights, score, next_state)

    return results


if __name__ == '__main__':
    # Load data
    schedule_df = pd.read_csv('../data/schedule_easy_without_experiment.CSV')
    priority_df = pd.read_csv('../data/priority.CSV')

    # Initialize schedule tensor
    weeks, days_per_week, classes_per_day = 21, 8, 13
    schedule_tensor = torch.zeros(weeks, days_per_week, classes_per_day)

    # Fill schedule tensor
    for time_slot in range(classes_per_day - 1):
        for day_index in range(1, days_per_week):
            value = schedule_df.iloc[time_slot].iloc[day_index]

            if value != 0 and value != "0":
                weeks_with_class = ([int(w) for w in value.split(', ')]
                                    if isinstance(value, str) else [value])

                for week in weeks_with_class:
                    if week < weeks:
                        schedule_tensor[week, day_index, time_slot + 1] = 1

    # Evaluate schedules
    results = judge_schedule(schedule_tensor, priority_df)
    print(f"Schedule scores: {results}")
